import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.gridspec import GridSpec
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import talib
import warnings
from scipy import stats

warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 初始化pro接口
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')


# 获取股票基本信息
def get_stock_basic():
    df = pro.stock_basic(**{
        "ts_code": "",
        "name": "",
        "exchange": "",
        "market": "",
        "is_hs": "",
        "list_status": "L",
        "limit": "",
        "offset": ""
    }, fields=[
        "ts_code",
        "symbol",
        "name",
        "area",
        "industry",
        "cnspell",
        "market",
        "list_date",
        "act_name",
        "act_ent_type"
    ])
    return df


# 获取日线数据
def get_daily_data(ts_code, start_date, end_date):
    df = pro.daily(**{
        "ts_code": ts_code,
        "trade_date": "",
        "start_date": start_date,
        "end_date": end_date,
        "offset": "",
        "limit": ""
    }, fields=[
        "ts_code",
        "trade_date",
        "open",
        "high",
        "low",
        "close",
        "pre_close",
        "change",
        "pct_chg",
        "vol",
        "amount"
    ])
    # 按日期升序排列
    df = df.sort_values('trade_date')
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df.set_index('trade_date', inplace=True)
    return df


# 特征工程函数
def create_features(df, lookback=20):
    """创建用于机器学习的特征"""
    features_df = df.copy()

    # 价格相关特征
    features_df['returns'] = features_df['close'].pct_change()
    features_df['log_returns'] = np.log(1 + features_df['returns'])

    # 移动平均
    for window in [5, 10, 20, 60]:
        features_df[f'sma_{window}'] = features_df['close'].rolling(window=window).mean()
        features_df[f'sma_{window}_ratio'] = features_df['close'] / features_df[f'sma_{window}']

        features_df[f'ema_{window}'] = features_df['close'].ewm(span=window).mean()
        features_df[f'ema_{window}_ratio'] = features_df['close'] / features_df[f'ema_{window}']

    # 波动率特征
    features_df['volatility_5'] = features_df['returns'].rolling(window=5).std()
    features_df['volatility_20'] = features_df['returns'].rolling(window=20).std()

    # 交易量特征
    features_df['volume_ma_5'] = features_df['vol'].rolling(window=5).mean()
    features_df['volume_ma_20'] = features_df['vol'].rolling(window=20).mean()
    features_df['volume_ratio'] = features_df['vol'] / features_df['volume_ma_20']

    # 价格动量特征
    features_df['momentum_5'] = features_df['close'].pct_change(periods=5)
    features_df['momentum_10'] = features_df['close'].pct_change(periods=10)
    features_df['momentum_20'] = features_df['close'].pct_change(periods=20)

    # 计算技术指标
    # MACD
    features_df['macd'], features_df['macd_signal'], features_df['macd_hist'] = talib.MACD(
        features_df['close'], fastperiod=12, slowperiod=26, signalperiod=9)

    # RSI
    features_df['rsi_14'] = talib.RSI(features_df['close'], timeperiod=14)

    # 布林带
    features_df['upper_band'], features_df['middle_band'], features_df['lower_band'] = talib.BBANDS(
        features_df['close'], timeperiod=20)
    features_df['bb_width'] = (features_df['upper_band'] - features_df['lower_band']) / features_df['middle_band']

    # KDJ指标
    features_df['k'], features_df['d'] = talib.STOCH(
        features_df['high'], features_df['low'], features_df['close'],
        fastk_period=9, slowk_period=3, slowk_matype=0, slowd_period=3, slowd_matype=0)
    features_df['j'] = 3 * features_df['k'] - 2 * features_df['d']

    # 计算未来N天的收益率作为标签
    future_return_days = 5
    features_df['future_return'] = features_df['close'].pct_change(future_return_days).shift(-future_return_days)

    # 创建分类标签：1表示上涨，0表示下跌
    features_df['target'] = (features_df['future_return'] > 0).astype(int)

    # 删除NaN值
    features_df = features_df.dropna()

    return features_df


# 交易信号生成器
class SignalGenerator:
    def __init__(self, strategy_type='ml', params=None):
        """
        初始化信号生成器

        参数:
            strategy_type: 策略类型，'ml'表示机器学习策略，'kdj'表示KDJ策略
            params: 策略参数
        """
        self.strategy_type = strategy_type
        self.params = params or {}

        # 设置默认参数
        if strategy_type == 'ml':
            self.default_params = {
                'n_estimators': 100,
                'lookback': 20,
                'threshold': 0.5  # 买入信号的概率阈值
            }
        elif strategy_type == 'kdj':
            self.default_params = {
                'n': 9,
                'm1': 3,
                'm2': 3
            }
        else:
            raise ValueError(f"不支持的策略类型: {strategy_type}")

        # 合并用户参数和默认参数
        self.params = {**self.default_params, **self.params}

        # 初始化模型
        if strategy_type == 'ml':
            self.model = RandomForestClassifier(n_estimators=self.params['n_estimators'], random_state=42)
            self.scaler = StandardScaler()

    def train(self, df):
        """训练模型（仅适用于机器学习策略）"""
        if self.strategy_type != 'ml':
            return

        # 创建特征
        features_df = create_features(df, lookback=self.params['lookback'])

        # 选择特征和目标变量
        feature_columns = [
            'sma_5_ratio', 'sma_10_ratio', 'sma_20_ratio', 'sma_60_ratio',
            'ema_5_ratio', 'ema_10_ratio', 'ema_20_ratio', 'ema_60_ratio',
            'volatility_5', 'volatility_20',
            'volume_ratio',
            'momentum_5', 'momentum_10', 'momentum_20',
            'macd', 'macd_signal', 'macd_hist',
            'rsi_14',
            'bb_width',
            'k', 'd', 'j'
        ]

        X = features_df[feature_columns]
        y = features_df['target']

        # 标准化特征
        X_scaled = self.scaler.fit_transform(X)

        # 训练模型
        self.model.fit(X_scaled, y)

        # 保存特征列，用于预测
        self.feature_columns = feature_columns

        # 保存最后N天的数据，用于实时预测
        self.lookback_data = features_df.iloc[-self.params['lookback']:]

        return self.model

    def generate_signals(self, df):
        """生成交易信号"""
        if self.strategy_type == 'ml':
            return self._generate_ml_signals(df)
        elif self.strategy_type == 'kdj':
            return self._generate_kdj_signals(df)

    def _generate_ml_signals(self, df):
        """生成机器学习策略的交易信号"""
        # 创建特征
        features_df = create_features(df, lookback=self.params['lookback'])

        # 选择特征
        X = features_df[self.feature_columns]

        # 标准化特征
        X_scaled = self.scaler.transform(X)

        # 预测
        y_proba = self.model.predict_proba(X_scaled)[:, 1]

        # 生成信号
        signals = pd.Series(0, index=features_df.index)
        signals[y_proba > self.params['threshold']] = 1  # 买入信号
        signals[y_proba < (1 - self.params['threshold'])] = -1  # 卖出信号

        # 合并回原始数据
        df_with_signals = df.copy()
        df_with_signals['signal'] = signals

        return df_with_signals

    def _generate_kdj_signals(self, df):
        """生成KDJ策略的交易信号"""
        # 计算KDJ指标
        df['low_n'] = df['low'].rolling(self.params['n']).min()
        df['high_n'] = df['high'].rolling(self.params['n']).max()

        # 计算RSV
        df['rsv'] = 100 * (df['close'] - df['low_n']) / (df['high_n'] - df['low_n'])

        # 计算K、D、J值
        df['k'] = df['rsv'].ewm(com=self.params['m1'] - 1, adjust=False).mean()
        df['d'] = df['k'].ewm(com=self.params['m2'] - 1, adjust=False).mean()
        df['j'] = 3 * df['k'] - 2 * df['d']

        # 生成交易信号
        df['signal'] = 0

        # KDJ金叉: K线上穿D线，买入信号
        df.loc[(df['k'] > df['d']) & (df['k'].shift(1) <= df['d'].shift(1)), 'signal'] = 1

        # KDJ死叉: K线下穿D线，卖出信号
        df.loc[(df['k'] < df['d']) & (df['k'].shift(1) >= df['d'].shift(1)), 'signal'] = -1

        return df


# 止损止盈管理器
class StopManager:
    def __init__(self, stop_loss_pct=0.05, take_profit_pct=0.10, trailing_stop_pct=0.03):
        """
        初始化止损止盈管理器

        参数:
            stop_loss_pct: 固定止损百分比
            take_profit_pct: 固定止盈百分比
            trailing_stop_pct: 追踪止损百分比
        """
        self.stop_loss_pct = stop_loss_pct
        self.take_profit_pct = take_profit_pct
        self.trailing_stop_pct = trailing_stop_pct

    def apply_stops(self, df):
        """应用止损止盈规则"""
        df_with_stops = df.copy()
        df_with_stops['position'] = 0  # 0表示空仓，1表示持有

        # 遍历每个交易日
        for i in range(1, len(df_with_stops)):
            # 获取前一天的持仓状态
            prev_position = df_with_stops['position'].iloc[i - 1]

            # 如果前一天为空仓
            if prev_position == 0:
                # 如果当天有买入信号
                if df_with_stops['signal'].iloc[i] == 1:
                    df_with_stops.loc[df_with_stops.index[i], 'position'] = 1
                    # 设置止损价和止盈价
                    entry_price = df_with_stops['close'].iloc[i]
                    df_with_stops.loc[df_with_stops.index[i], 'stop_loss'] = entry_price * (1 - self.stop_loss_pct)
                    df_with_stops.loc[df_with_stops.index[i], 'take_profit'] = entry_price * (1 + self.take_profit_pct)
                    df_with_stops.loc[df_with_stops.index[i], 'trailing_stop'] = entry_price * (
                                1 - self.trailing_stop_pct)
                    df_with_stops.loc[df_with_stops.index[i], 'highest_price'] = entry_price
                else:
                    df_with_stops.loc[df_with_stops.index[i], 'position'] = 0
            # 如果前一天持有仓位
            else:
                # 获取当前价格和前一天的止损止盈价
                current_price = df_with_stops['close'].iloc[i]
                prev_stop_loss = df_with_stops['stop_loss'].iloc[i - 1]
                prev_take_profit = df_with_stops['take_profit'].iloc[i - 1]
                prev_trailing_stop = df_with_stops['trailing_stop'].iloc[i - 1]
                prev_highest_price = df_with_stops['highest_price'].iloc[i - 1]

                # 更新最高价
                highest_price = max(prev_highest_price, current_price)

                # 更新追踪止损价
                trailing_stop = highest_price * (1 - self.trailing_stop_pct)

                # 检查是否触发止损或止盈
                if current_price <= min(prev_stop_loss, trailing_stop) or current_price >= prev_take_profit:
                    # 触发止损或止盈，平仓
                    df_with_stops.loc[df_with_stops.index[i], 'position'] = 0
                    # 记录交易类型
                    if current_price <= min(prev_stop_loss, trailing_stop):
                        df_with_stops.loc[df_with_stops.index[i], 'trade_type'] = 'stop_loss'
                    else:
                        df_with_stops.loc[df_with_stops.index[i], 'trade_type'] = 'take_profit'
                # 如果当天有卖出信号，也平仓
                elif df_with_stops['signal'].iloc[i] == -1:
                    df_with_stops.loc[df_with_stops.index[i], 'position'] = 0
                    df_with_stops.loc[df_with_stops.index[i], 'trade_type'] = 'signal_exit'
                else:
                    # 继续持有仓位
                    df_with_stops.loc[df_with_stops.index[i], 'position'] = 1
                    df_with_stops.loc[df_with_stops.index[i], 'stop_loss'] = prev_stop_loss
                    df_with_stops.loc[df_with_stops.index[i], 'take_profit'] = prev_take_profit
                    df_with_stops.loc[df_with_stops.index[i], 'trailing_stop'] = trailing_stop
                    df_with_stops.loc[df_with_stops.index[i], 'highest_price'] = highest_price

        return df_with_stops


# 回测引擎
class BacktestEngine:
    def __init__(self, initial_capital=100000.0, commission_rate=0.0003, slippage=0.0001):
        """
        初始化回测引擎

        参数:
            initial_capital: 初始资金
            commission_rate: 交易费率
            slippage: 滑点
        """
        self.initial_capital = initial_capital
        self.commission_rate = commission_rate
        self.slippage = slippage

    def run_backtest(self, df_with_positions):
        """运行回测"""
        df = df_with_positions.copy()

        # 初始化资金和持仓
        df['capital'] = self.initial_capital
        df['shares'] = 0
        df['position_value'] = 0
        df['total_equity'] = self.initial_capital

        # 交易记录
        trades = []

        # 遍历每个交易日
        for i in range(1, len(df)):
            # 获取前一天的状态
            prev_capital = df['capital'].iloc[i - 1]
            prev_shares = df['shares'].iloc[i - 1]
            current_price = df['close'].iloc[i]
            current_position = df['position'].iloc[i]

            # 如果前一天为空仓，今天要买入
            if prev_shares == 0 and current_position == 1:
                # 计算可买入的股数（假设全仓买入）
                commission = prev_capital * self.commission_rate
                slippage_cost = prev_capital * self.slippage
                available_capital = prev_capital - commission - slippage_cost
                shares_to_buy = int(available_capital / current_price / 100) * 100  # 取整百

                if shares_to_buy > 0:
                    cost = shares_to_buy * current_price
                    actual_commission = cost * self.commission_rate
                    actual_slippage = cost * self.slippage
                    total_cost = cost + actual_commission + actual_slippage

                    # 更新状态
                    df.loc[df.index[i], 'shares'] = shares_to_buy
                    df.loc[df.index[i], 'capital'] = prev_capital - total_cost
                    df.loc[df.index[i], 'position_value'] = shares_to_buy * current_price
                    df.loc[df.index[i], 'total_equity'] = df.loc[df.index[i], 'capital'] + df.loc[
                        df.index[i], 'position_value']

                    # 记录交易
                    trades.append({
                        'date': df.index[i],
                        'type': 'buy',
                        'price': current_price,
                        'shares': shares_to_buy,
                        'cost': total_cost,
                        'commission': actual_commission,
                        'slippage': actual_slippage
                    })
                else:
                    df.loc[df.index[i], 'shares'] = 0
                    df.loc[df.index[i], 'capital'] = prev_capital
                    df.loc[df.index[i], 'position_value'] = 0
                    df.loc[df.index[i], 'total_equity'] = prev_capital

            # 如果前一天持有仓位，今天要卖出
            elif prev_shares > 0 and current_position == 0:
                # 计算卖出收益
                revenue = prev_shares * current_price
                commission = revenue * self.commission_rate
                slippage = revenue * self.slippage
                net_revenue = revenue - commission - slippage

                # 更新状态
                df.loc[df.index[i], 'shares'] = 0
                df.loc[df.index[i], 'capital'] = prev_capital + net_revenue
                df.loc[df.index[i], 'position_value'] = 0
                df.loc[df.index[i], 'total_equity'] = df.loc[df.index[i], 'capital']

                # 记录交易
                trades.append({
                    'date': df.index[i],
                    'type': 'sell',
                    'price': current_price,
                    'shares': prev_shares,
                    'revenue': net_revenue,
                    'commission': commission,
                    'slippage': slippage
                })

                # 记录交易类型
                if 'trade_type' in df.columns and pd.notna(df['trade_type'].iloc[i]):
                    trades[-1]['exit_reason'] = df['trade_type'].iloc[i]

            # 如果持仓状态不变
            else:
                df.loc[df.index[i], 'shares'] = prev_shares
                df.loc[df.index[i], 'capital'] = prev_capital
                df.loc[df.index[i], 'position_value'] = prev_shares * current_price
                df.loc[df.index[i], 'total_equity'] = df.loc[df.index[i], 'capital'] + df.loc[
                    df.index[i], 'position_value']

        # 计算收益率
        df['daily_return'] = df['total_equity'].pct_change()
        df['cumulative_return'] = (1 + df['daily_return']).cumprod() - 1

        # 计算买入持有收益率
        df['buy_hold_return'] = df['close'] / df['close'].iloc[0] - 1

        # 计算回撤
        df['peak'] = df['total_equity'].cummax()
        df['drawdown'] = (df['total_equity'] - df['peak']) / df['peak']

        # 计算绩效指标
        performance = self._calculate_performance(df, trades)

        return df, trades, performance

    def _calculate_performance(self, df, trades):
        """计算绩效指标"""
        performance = {}

        # 基本信息
        performance['start_date'] = df.index[0]
        performance['end_date'] = df.index[-1]
        performance['days'] = (df.index[-1] - df.index[0]).days

        # 收益率
        performance['total_return'] = df['cumulative_return'].iloc[-1] * 100
        performance['buy_hold_return'] = df['buy_hold_return'].iloc[-1] * 100

        # 年化收益率
        performance['annual_return'] = ((1 + performance['total_return'] / 100) ** (
                    365 / performance['days']) - 1) * 100
        performance['buy_hold_annual_return'] = ((1 + performance['buy_hold_return'] / 100) ** (
                    365 / performance['days']) - 1) * 100

        # 最大回撤
        performance['max_drawdown'] = df['drawdown'].min() * 100

        # 夏普比率
        risk_free_rate = 0.03  # 假设无风险利率为3%
        daily_returns = df['daily_return'].dropna()
        performance['sharpe_ratio'] = (performance['annual_return'] / 100 - risk_free_rate) / (
                    daily_returns.std() * np.sqrt(252))

        # 交易统计
        if trades:
            buy_trades = [t for t in trades if t['type'] == 'buy']
            sell_trades = [t for t in trades if t['type'] == 'sell']

            performance['total_trades'] = len(sell_trades)
            performance['win_trades'] = sum(
                1 for t in sell_trades if t['revenue'] > t['shares'] * buy_trades[sell_trades.index(t)]['price'])
            performance['win_rate'] = performance['win_trades'] / len(sell_trades) * 100 if sell_trades else 0

            # 平均收益/亏损
            if sell_trades:
                trade_returns = [(t['revenue'] / (t['shares'] * buy_trades[sell_trades.index(t)]['price']) - 1) * 100
                                 for t in sell_trades]
                performance['avg_win'] = np.mean([r for r in trade_returns if r > 0]) if any(
                    r > 0 for r in trade_returns) else 0
                performance['avg_loss'] = np.mean([r for r in trade_returns if r < 0]) if any(
                    r < 0 for r in trade_returns) else 0
                performance['profit_factor'] = abs(
                    sum(r for r in trade_returns if r > 0) / sum(r for r in trade_returns if r < 0)) if any(
                    r < 0 for r in trade_returns) else float('inf')

                # 最大盈利/亏损交易
                performance['max_win'] = max(trade_returns) if trade_returns else 0
                performance['max_loss'] = min(trade_returns) if trade_returns else 0

        # 超额收益
        performance['excess_return'] = performance['annual_return'] - performance['buy_hold_annual_return']

        return performance


# 可视化回测结果
def visualize_backtest(stock_name, backtest_df, trades, performance):
    """可视化回测结果"""
    # 创建一个带有多个子图的画布
    fig = plt.figure(figsize=(20, 20))
    gs = GridSpec(5, 1, figure=fig)

    # 子图1: 价格和交易信号
    ax1 = fig.add_subplot(gs[0, 0])
    ax1.plot(backtest_df.index, backtest_df['close'], label='收盘价', color='blue')

    # 标记买入和卖出信号
    buy_trades = [t for t in trades if t['type'] == 'buy']
    sell_trades = [t for t in trades if t['type'] == 'sell']

    buy_dates = [t['date'] for t in buy_trades]
    buy_prices = [t['price'] for t in buy_trades]

    sell_dates = [t['date'] for t in sell_trades]
    sell_prices = [t['price'] for t in sell_trades]

    ax1.scatter(buy_dates, buy_prices, marker='^', color='green', s=100, label='买入')
    ax1.scatter(sell_dates, sell_prices, marker='v', color='red', s=100, label='卖出')

    ax1.set_title(f'{stock_name} 价格和交易信号', fontsize=15)
    ax1.set_ylabel('价格', fontsize=12)
    ax1.grid(True, linestyle='--', alpha=0.7)
    ax1.legend()

    # 子图2: 策略权益曲线
    ax2 = fig.add_subplot(gs[1, 0])
    ax2.plot(backtest_df.index, backtest_df['total_equity'], label='策略权益', color='blue', linewidth=2)
    ax2.set_title('策略权益曲线', fontsize=15)
    ax2.set_ylabel('权益', fontsize=12)
    ax2.grid(True, linestyle='--', alpha=0.7)
    ax2.legend()

    # 子图3: 策略收益vs买入持有
    ax3 = fig.add_subplot(gs[2, 0])
    ax3.plot(backtest_df.index, backtest_df['cumulative_return'] * 100, label='策略收益(%)', color='blue', linewidth=2)
    ax3.plot(backtest_df.index, backtest_df['buy_hold_return'] * 100, label='买入持有收益(%)', color='red', linewidth=2)
    ax3.set_title('策略收益vs买入持有收益', fontsize=15)
    ax3.set_ylabel('收益率(%)', fontsize=12)
    ax3.grid(True, linestyle='--', alpha=0.7)
    ax3.legend()

    # 子图4: 回撤
    ax4 = fig.add_subplot(gs[3, 0])
    ax4.fill_between(backtest_df.index, backtest_df['drawdown'] * 100, 0, color='red', alpha=0.3)
    ax4.plot(backtest_df.index, backtest_df['drawdown'] * 100, color='red', linewidth=1)
    ax4.set_title('回撤', fontsize=15)
    ax4.set_ylabel('回撤(%)', fontsize=12)
    ax4.grid(True, linestyle='--', alpha=0.7)

    # 添加最大回撤信息
    max_drawdown = performance['max_drawdown']
    max_drawdown_date = backtest_df['drawdown'].idxmin()
    ax4.text(0.05, 0.95, f'最大回撤: {max_drawdown:.2f}%', transform=ax4.transAxes,
             fontsize=12, verticalalignment='top', bbox=dict(boxstyle='round', facecolor='red', alpha=0.2))

    # 子图5: 绩效指标雷达图
    ax5 = fig.add_subplot(gs[4, 0], polar=True)

    # 选择要显示的指标
    radar_metrics = ['年化收益率(%)', '最大回撤(%)', '夏普比率', '胜率(%)', '盈亏比']

    # 转换指标值为0-1之间
    max_values = {
        '年化收益率(%)': max(performance['annual_return'], performance['buy_hold_annual_return'], 20),
        '最大回撤(%)': 30,  # 负值，越小越好
        '夏普比率': max(performance['sharpe_ratio'], 2),
        '胜率(%)': 100,
        '盈亏比': max(performance.get('profit_factor', 0), 3)
    }

    # 准备雷达图数据
    strategy_radar = [
        performance['annual_return'] / max_values['年化收益率(%)'],
        1 - abs(performance['max_drawdown'] / max_values['最大回撤(%)']),  # 最大回撤取反，因为越小越好
        performance['sharpe_ratio'] / max_values['夏普比率'],
        performance.get('win_rate', 0) / max_values['胜率(%)'],
        performance.get('profit_factor', 0) / max_values['盈亏比']
    ]

    buy_hold_radar = [
        performance['buy_hold_annual_return'] / max_values['年化收益率(%)'],
        1 - abs(performance['max_drawdown'] / max_values['最大回撤(%)']),  # 最大回撤取反，因为越小越好
        0,  # 买入持有没有夏普比率
        0,  # 买入持有没有胜率
        0  # 买入持有没有盈亏比
    ]

    # 角度
    angles = np.linspace(0, 2 * np.pi, len(radar_metrics), endpoint=False).tolist()
    angles += angles[:1]  # 闭合雷达图

    # 添加数据点
    strategy_radar += strategy_radar[:1]
    buy_hold_radar += buy_hold_radar[:1]

    # 绘制雷达图
    ax5.plot(angles, strategy_radar, 'o-', linewidth=2, label='策略绩效', color='blue')
    ax5.plot(angles, buy_hold_radar, 'o-', linewidth=2, label='买入持有', color='red')
    ax5.fill(angles, strategy_radar, alpha=0.25, color='blue')
    ax5.fill(angles, buy_hold_radar, alpha=0.25, color='red')

    # 设置标签
    ax5.set_thetagrids(np.degrees(angles[:-1]), radar_metrics)
    ax5.set_title('策略绩效雷达图', fontsize=15)
    ax5.legend(loc='upper right', bbox_to_anchor=(0.1, 0.1))

    plt.tight_layout()
    plt.savefig(f'{stock_name}_回测结果.png', dpi=300, bbox_inches='tight')
    plt.show()


# 主函数
def main():
    # 获取股票基本信息
    stocks = get_stock_basic()

    # 选择几支股票进行分析（选择流动性好的大盘股）
    selected_stocks = stocks[(stocks['industry'].isin(['银行', '白酒', '电子', '医药'])) &
                             (stocks['market'] == '主板') &
                             (stocks['list_date'] < '20150101')].head(5)

    # 设置时间范围
    start_date = '20180101'
    end_date = '20230101'

    # 回测结果汇总
    results_summary = []

    # 对每支股票进行回测
    for _, stock in selected_stocks.iterrows():
        ts_code = stock['ts_code']
        stock_name = stock['name']
        print(f"\n正在回测 {stock_name} ({ts_code})...")

        # 获取日线数据
        daily_data = get_daily_data(ts_code, start_date, end_date)

        # 创建信号生成器（使用机器学习策略）
        signal_generator = SignalGenerator(strategy_type='ml', params={
            'n_estimators': 100,
            'threshold': 0.55  # 提高买入阈值，减少交易频率
        })

        # 训练模型
        signal_generator.train(daily_data)

        # 生成交易信号
        df_with_signals = signal_generator.generate_signals(daily_data)

        # 应用止损止盈
        stop_manager = StopManager(
            stop_loss_pct=0.08,  # 8%止损
            take_profit_pct=0.15,  # 15%止盈
            trailing_stop_pct=0.05  # 5%追踪止损
        )

        df_with_positions = stop_manager.apply_stops(df_with_signals)

        # 运行回测
        backtest_engine = BacktestEngine(
            initial_capital=100000.0,
            commission_rate=0.0003,  # 万三佣金
            slippage=0.0001  # 千分之一滑点
        )

        backtest_df, trades, performance = backtest_engine.run_backtest(df_with_positions)

        # 保存结果
        results_summary.append({
            'stock_name': stock_name,
            'ts_code': ts_code,
            'total_return': performance['total_return'],
            'annual_return': performance['annual_return'],
            'max_drawdown': performance['max_drawdown'],
            'sharpe_ratio': performance['sharpe_ratio'],
            'win_rate': performance.get('win_rate', 0),
            'profit_factor': performance.get('profit_factor', 0),
            'excess_return': performance['excess_return']
        })

        # 可视化回测结果
        visualize_backtest(stock_name, backtest_df, trades, performance)

        # 打印回测绩效
        print(f"\n{stock_name} 回测绩效:")
        for metric, value in performance.items():
            print(f"{metric}: {value:.4f}")

        # 打印交易记录摘要
        print(f"\n交易记录摘要:")
        print(f"总交易次数: {len(trades) // 2}")
        if trades:
            buy_trades = [t for t in trades if t['type'] == 'buy']
            sell_trades = [t for t in trades if t['type'] == 'sell']

            if sell_trades:
                win_trades = sum(
                    1 for t in sell_trades if t['revenue'] > t['shares'] * buy_trades[sell_trades.index(t)]['price'])
                win_rate = win_trades / len(sell_trades) * 100 if sell_trades else 0

                print(f"盈利交易: {win_trades} ({win_rate:.2f}%)")
                print(f"亏损交易: {len(sell_trades) - win_trades} ({100 - win_rate:.2f}%)")

                # 统计不同退出原因的交易数量
                exit_reasons = {}
                for t in sell_trades:
                    reason = t.get('exit_reason', 'signal')
                    exit_reasons[reason] = exit_reasons.get(reason, 0) + 1

                print("\n退出原因统计:")
                for reason, count in exit_reasons.items():
                    print(f"{reason}: {count} ({count / len(sell_trades) * 100:.2f}%)")

        print("-" * 70)

    # 汇总结果
    if results_summary:
        print("\n回测结果汇总:")
        results_df = pd.DataFrame(results_summary)
        print(results_df)

        # 计算平均绩效
        print("\n平均绩效:")
        print(f"平均年化收益率: {results_df['annual_return'].mean():.2f}%")
        print(f"平均最大回撤: {results_df['max_drawdown'].mean():.2f}%")
        print(f"平均夏普比率: {results_df['sharpe_ratio'].mean():.2f}")
        print(f"平均胜率: {results_df['win_rate'].mean():.2f}%")
        print(f"平均盈亏比: {results_df['profit_factor'].mean():.2f}")

        # 可视化汇总结果
        plt.figure(figsize=(15, 10))

        # 年化收益率对比
        plt.subplot(2, 2, 1)
        sns.barplot(x='stock_name', y='annual_return', data=results_df)
        plt.title('年化收益率对比')
        plt.xlabel('股票')
        plt.ylabel('年化收益率(%)')
        plt.xticks(rotation=45)

        # 最大回撤对比
        plt.subplot(2, 2, 2)
        sns.barplot(x='stock_name', y='max_drawdown', data=results_df)
        plt.title('最大回撤对比')
        plt.xlabel('股票')
        plt.ylabel('最大回撤(%)')
        plt.xticks(rotation=45)

        # 夏普比率对比
        plt.subplot(2, 2, 3)
        sns.barplot(x='stock_name', y='sharpe_ratio', data=results_df)
        plt.title('夏普比率对比')
        plt.xlabel('股票')
        plt.ylabel('夏普比率')
        plt.xticks(rotation=45)

        # 胜率对比
        plt.subplot(2, 2, 4)
        sns.barplot(x='stock_name', y='win_rate', data=results_df)
        plt.title('胜率对比')
        plt.xlabel('股票')
        plt.ylabel('胜率(%)')
        plt.xticks(rotation=45)

        plt.tight_layout()
        plt.savefig('回测结果汇总.png', dpi=300, bbox_inches='tight')
        plt.show()


if __name__ == "__main__":
    main()
